import random
import torch, math
from cell_genetics import CellGenetics
from enum import Enum, auto


class CellDataIndex(Enum):
    ALIVE = auto()
    ENERGY = auto()
    AGE = auto()
    LIFETIME = auto()
    DIVISION_THRESHOLD = auto()
    LOW_METABOLISM_HIGH_METABOLISM = auto()
    MUTATION_RATE = auto()
    AUTOTROPHIC_HETEROTROPHIC = auto()  # 新增：自养标志
    DEFENSE_OFFENSE = auto()  # 新增：防御/进攻参数
    MOVE_ENERGY_COST = auto()  # 新增：移动能量消耗
    STATIC_MOTION = auto()  # 新增：移动/静止
    SIZE = auto()  # 新增：细胞大小
    ATTRACT_REPEL = auto()  # 新增：吸引/排斥参数


class Cell:
    def __init__(self, genetic_code=None, grid=None):
        """
        初始化细胞对象，设置能量、年龄、毒素、寿命和遗传特征。
        :param genetic_code: 遗传代码，默认为None，会随机生成。
        :param grid: 网格对象，默认为None。
        """
        self.genetics = CellGenetics(genetic_code)
        self.energy = 500  # 能量初始值固定为500
        self.age = 0
        self.lifetime = 0
        # 设置细胞坐标属性
        self.x = 0
        self.y = 0
        self.grid = grid
        if self.grid:
            self.grid.grid[self.x][self.y] = self
        self._initialize_traits()  # @call _initialize_traits

    def _initialize_traits(self):
        """
        从遗传代码中解码并初始化细胞的遗传特征。
        调用CellGenetics.decode_gene方法，获取包括突变率、寿命、分裂阈值、代谢率、自养标志、能量来源类型和防御/进攻参数等特征。
        """
        traits = self.genetics.decode_gene()
        self.mutation_rate = traits['mutation_rate']['value']
        self.long_life_short_life = traits['long_life_short_life']['value']
        self.division_threshold = traits['division_threshold']['value']
        self.low_metabolism_high_metabolism = traits['low_metabolism_high_metabolism']['value']
        self.autotrophic_heterotrophic = traits['autotrophic_heterotrophic']['value']
        self.defense_offense = traits['defense_offense']['value']
        self.move_energy_cost = traits['move_energy_cost']['value']
        self.static_motion = traits['static_motion']['value']  # 新增：移动/静止
        self.current_energy = self.energy  # 添加实时能量参数

    def metabolize(self, env):
        """
        执行细胞的代谢过程，包括能量消耗和状态检查。
        :param env: 环境参数，当前未使用。
        :return: 细胞的新状态，可能为'divide'、'die'或'survive'。
        """
        # 独立更新能量，不受突变影响
        if self.autotrophic_heterotrophic:
            # 自养细胞能量增加
            self.energy += self.autotrophic_heterotrophic
        else:
            # 异养细胞能量消耗，实现非线性衰减
            if self.low_metabolism_high_metabolism <= 16:
                self.energy -= self.low_metabolism_high_metabolism * 0.8
            else:
                self.energy -= 16 * 0.8 + (self.low_metabolism_high_metabolism - 16) * 0.2  # 修改：降低高代谢细胞的能量惩罚

        self.age += 1

        # 分裂条件（动态阈值）
        reproduction_threshold = max(self.low_metabolism_high_metabolism * 1.2, 0.3 * 10)
        if self.energy > reproduction_threshold:
            return 'divide'

        # 死亡条件
        if self.age > self.long_life_short_life * 10 * (1 - self.mutation_rate * 2):
            return 'die'

        return 'survive'

    def move(self, x, y, dx, dy):
        """
        细胞移动逻辑。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        :param dx: 移动的x方向。
        :param dy: 移动的y方向。
        :return: 是否成功移动。
        """
        if not hasattr(self, 'grid') or not self.grid:
            return False
        neighbors = self.grid.get_neighbors(x, y)  # 获取邻居细胞
        # 计算基因相似度并调整能耗
        genetic_similarities = [self.genetic_similarity(n) for n in neighbors]
        if genetic_similarities:
            max_similarity = max(genetic_similarities)
            # 新移动消耗公式：⌊静止-运动参数/10⌋ × √当前能量
            cost = math.floor(self.static_motion / 10) * math.sqrt(max(0, self.energy)) * (1 - 0.5 * max_similarity)
        else:
            cost = math.floor(self.static_motion / 10) * math.sqrt(max(0, self.energy))

        # 移动步长为static_motion参数值
        distance = self.static_motion
        # 计算曼哈顿距离
        manhattan_distance = abs(dx) + abs(dy)
        if manhattan_distance > distance:
            return False  # 曼哈顿距离超过限制，无法移动
        if self.energy < cost:
            return False  # 能量不足，无法移动

        new_x, new_y = x + dx * distance, y + dy * distance
        if 0 <= new_x < self.grid.width and 0 <= new_y < self.grid.height:
            try:
                if not self.grid.grid[new_x][new_y]:
                    self.grid.grid[new_x][new_y] = self
                    self.grid.grid[x][y] = None
                    self.x = new_x
                    self.y = new_y
            except IndexError:
                return False
        else:
            return False

        # 异养细胞移动优势：降低能量消耗20%，增加移动步长50%
        if not self.autotrophic_heterotrophic:
            self.energy -= cost * manhattan_distance * 0.8
        else:
            self.energy -= cost * manhattan_distance
        return True

    def genetic_similarity(self, other):
        """
        计算两个细胞的基因相似度。
        :param other: 另一个细胞。
        :return: 基因相似度。
        """
        return sum(1 for a, b in zip(self.genetics.genetic_code, other.genetics.genetic_code) if a == b) / len(self.genetics.genetic_code)


class SimulationEngine:
    def __init__(self, width=100, height=100, total_generations=100):
        """
        初始化模拟引擎，设置网格大小、设备、细胞数据、环境参数和初始状态。
        :param width: 网格宽度。
        :param height: 网格高度。
        :param total_generations: 总代数。
        """
        self.width = width
        self.height = height
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.cell_data = torch.zeros((height, width, len(CellDataIndex)), dtype=torch.float32, device=self.device)  # [alive, energy, age, lifetime, division_threshold, metabolism, autotrophic_heterotrophic, defense_offense, move_energy_cost, static_motion]
        self.grid = [[None for _ in range(width)] for _ in range(height)]
        self.env_params = {
            'mutation_rate': 0.001,
            'nutrient_distribution': 'uniform',
        }
        self.gene_stats = {}
        self.initial_state = None  # 用于存储初始状态
        self.current_generation = 0
        self.total_generations = 0
        self._save_initial_state()  # @call _save_initial_state
        self._grid_initialization()  # @call _grid_initialization

    def _grid_initialization(self):
        """
        初始化网格，随机放置细胞。
        """
        for i in range(self.height):
            for j in range(self.width):
                if random.random() < 0.1:  # 10%的概率初始化一个细胞
                    cell = Cell()
                    cell.grid = self  # 设置grid属性
                    self.grid[i][j] = cell
                    self.cell_data[i, j, CellDataIndex.ALIVE.value] = 1  # 标记为存活
                    self.cell_data[i, j, CellDataIndex.ENERGY.value] = cell.energy
                    self.cell_data[i, j, CellDataIndex.AGE.value] = cell.age
                    self.cell_data[i, j, CellDataIndex.LIFETIME.value] = cell.lifetime
                    self.cell_data[i, j, CellDataIndex.DIVISION_THRESHOLD.value] = cell.division_threshold
                    self.cell_data[i, j, CellDataIndex.LOW_METABOLISM_HIGH_METABOLISM.value] = cell.low_metabolism_high_metabolism
                    self.cell_data[i, j, CellDataIndex.AUTOTROPHIC_HETEROTROPHIC.value] = cell.autotrophic_heterotrophic  # 新增：自养标志
                    self.cell_data[i, j, CellDataIndex.DEFENSE_OFFENSE.value] = cell.defense_offense  # 新增：防御/进攻参数
                    self.cell_data[i, j, CellDataIndex.MOVE_ENERGY_COST.value] = cell.move_energy_cost  # 新增：移动能量消耗
                    self.cell_data[i, j, CellDataIndex.STATIC_MOTION.value] = cell.static_motion if hasattr(cell, 'static_motion') else 0  # 新增：移动/静止

    def get_neighbors(self, x, y):
        """
        获取指定位置周围的邻居细胞。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        :return: 邻居细胞列表。
        """
        neighbors = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                if dx == 0 and dy == 0:
                    continue
                nx, ny = x + dx, y + dy
                if 0 <= nx < len(self.grid) and 0 <= ny < len(self.grid[0]):
                    if self.grid[nx][ny]:
                        neighbors.append(self.grid[nx][ny])
        return neighbors

    def _get_min_division_threshold(self, x, y):
        """
        获取周围同种细胞的最小分裂阈值。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        :return: 最小分裂阈值，若没有邻居则返回无穷大。
        """
        neighbors = self.get_neighbors(x, y)
        if not neighbors:
            return float('inf')
        return min(neighbor.division_threshold for neighbor in neighbors)

    def _get_cell_traits(self, x, y):
        """
        获取指定位置细胞的特征。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        :return: 细胞特征字典，若无细胞则返回None。
        """
        cell = self.grid[x][y]
        if not cell:
            return None
        return {
            'low_metabolism_high_metabolism': cell.low_metabolism_high_metabolism,
            'mutation_rate': cell.mutation_rate,
            'toxin_secretion': cell.toxin_secretion,
            'division_threshold': cell.division_threshold,
            'long_life_short_life': cell.long_life_short_life,
            'toxin_resistance': cell.toxin_resistance,
            'autotrophic_heterotrophic': cell.autotrophic_heterotrophic,  # 新增：自养标志
            'defense_offense': cell.defense_offense,  # 新增：防御/进攻参数
            'move_energy_cost': cell.move_energy_cost,  # 新增：移动能量消耗
            'static_motion': cell.static_motion  # 新增：移动/静止
        }

    def update_cell(self, x, y):
        """
        更新指定位置的细胞状态。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        """
        cell = self.grid[x][y]
        if not cell:
            return

        result = cell.metabolize(self.env_params)  # @call metabolize
        if result == 'divide':
            self._divide_cell(x, y)  # @call _divide_cell
        elif result == 'die':
            self.grid[x][y] = None
        else:
            # 能量掠夺逻辑
            self._energy_raiding(x, y)  # @call _energy_raiding
            # 能量扩散逻辑
            self._energy_diffusion(x, y)  # @call _energy_diffusion
            # 细胞移动逻辑
            self._cell_movement(x, y)  # @call _cell_movement

    def _divide_cell(self, x, y):
        """
        细胞分裂逻辑。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        """
        cell = self.grid[x][y]
        if not cell:
            return

        # 检查繁殖间隔条件
        if cell.age % cell.low_metabolism_high_metabolism != 0:
            return

        # 计算系统平均能量
        total_energy = torch.sum(self.cell_data[:, :, CellDataIndex.ENERGY.value])
        alive_count = torch.sum(self.cell_data[:, :, CellDataIndex.ALIVE.value])
        avg_energy = (total_energy / alive_count).item() if alive_count > 0 else 0
        
        # 计算繁殖阈值
        reproduction_threshold = max(cell.low_metabolism_high_metabolism * 1.2, 0.3 * avg_energy)
        
        # 检查能量是否达到繁殖阈值
        if cell.energy < reproduction_threshold:
            return

        # 寻找符合条件的空位
        directions = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]
        random.shuffle(directions)
        
        for dx, dy in directions:
            nx, ny = x + dx, y + dy
            if 0 <= nx < self.height and 0 <= ny < self.width and self.grid[nx][ny] is None:
                # 检查相邻格子能量是否大于等于代谢值的2倍
                if self.cell_data[nx, ny, CellDataIndex.ENERGY.value] >= cell.low_metabolism_high_metabolism * 2:
                    # 创建新细胞
                    new_cell = Cell(cell.genetics.mutate(cell.genetics.genetic_code, cell.mutation_rate))
                    new_cell.lifetime = cell.lifetime + 1
                    new_cell.energy = cell.energy / 2
                    new_cell.grid = self
                    new_cell._initialize_traits()
                    cell.energy /= 2
                    
                    # 放置新细胞
                    self.grid[nx][ny] = new_cell
                    self.cell_data[nx, ny, CellDataIndex.ALIVE.value] = 1
                    self.cell_data[nx, ny, CellDataIndex.ENERGY.value] = new_cell.energy
                    self.cell_data[nx, ny, CellDataIndex.AGE.value] = new_cell.age
                    self.cell_data[nx, ny, CellDataIndex.LIFETIME.value] = new_cell.lifetime
                    self.cell_data[nx, ny, CellDataIndex.LOW_METABOLISM_HIGH_METABOLISM.value] = new_cell.low_metabolism_high_metabolism
                    self.cell_data[nx, ny, CellDataIndex.MUTATION_RATE.value] = new_cell.mutation_rate
                    self.cell_data[nx, ny, CellDataIndex.AUTOTROPHIC_HETEROTROPHIC.value] = new_cell.autotrophic_heterotrophic
                    self.cell_data[nx, ny, CellDataIndex.DEFENSE_OFFENSE.value] = new_cell.defense_offense
                    self.cell_data[nx, ny, CellDataIndex.MOVE_ENERGY_COST.value] = new_cell.move_energy_cost
                    self.cell_data[nx, ny, CellDataIndex.STATIC_MOTION.value] = new_cell.static_motion
                    break  # 新增：移动/静止

    def _save_initial_state(self):
        """
        保存模拟的初始状态。
        """
        self.initial_state = {
            'cell_data': self.cell_data.clone(),
            'grid': [[cell.genetics.genetic_code if cell else None for cell in row] for row in self.grid],
            'env_params': self.env_params.copy(),
            'generation': 0  # 明确标记为第0代
        }

    def _update_gene_stats(self):
        """
        更新基因统计信息。
        """
        self.gene_stats = {}
        for row in self.grid:
            for cell in row:
                if cell:
                    code = cell.genetics.genetic_code
                    code_tuple = tuple(code)
                    if code_tuple not in self.gene_stats:
                        self.gene_stats[code_tuple] = {'count': 0, 'age': 0}
                    self.gene_stats[code_tuple]['count'] += 1
                    self.gene_stats[code_tuple]['age'] += cell.age

    def _energy_raiding(self, x, y):
        """
        能量掠夺机制，基于曼哈顿距离计算能量掠夺。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        """
        cell = self.grid[x][y]
        if not cell or cell.defense_offense <= 0:
            return

        # 扩大异养生物的侦测范围到7格
        neighbors = []
        for dx in range(-7, 8):
            for dy in range(-7, 8):
                if dx == 0 and dy == 0:
                    continue
                nx, ny = x + dx, y + dy
                if 0 <= nx < self.height and 0 <= ny < self.width and self.grid[nx][ny]:
                    neighbors.append(self.grid[nx][ny])
        
        if not neighbors:
            return

        # 优化能量转移算法，考虑距离衰减因子
        for neighbor in neighbors:
            distance = abs(neighbor.x - x) + abs(neighbor.y - y)
            if distance <= 2:
                # 基于防御-进攻参数的能量掠夺算法，确保差值大于0
                if cell.defense_offense - neighbor.defense_offense > 0:
                    # 进攻值=15的细胞吸收死亡能量时，额外获得20%能量加成
                    bonus = 1.2 if cell.defense_offense >= 15 else 1.0
                    # 距离衰减因子：1/distance
                    distance_factor = 1.0 / distance
                    energy_transfer = max(0, min(
                        (cell.defense_offense - neighbor.defense_offense) * bonus * distance_factor * 1.3,  # 提高掠夺效率30%
                        cell.energy,
                        1000
                    ))
                    if energy_transfer > 0:
                        cell.energy -= energy_transfer
                        neighbor.energy += energy_transfer

    def _energy_diffusion(self, x, y):
        """
        能量扩散机制，使用曼哈顿距离计算能量扩散。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        """
        cell = self.grid[x][y]
        if not cell:
            return

        neighbors = self.get_neighbors(x, y)
        if not neighbors:
            return

        # 计算曼哈顿距离
        for neighbor in neighbors:
            distance = abs(neighbor.x - x) + abs(neighbor.y - y)
            if distance <= 3:
                # 内圈扩散：曼哈顿距离≤3 → 每轮+⌈异养-自养参数/1.5⌉
                energy_transfer = math.ceil((neighbor.low_metabolism_high_metabolism - neighbor.autotrophic_heterotrophic) / 1.5)
            elif distance <= 6:
                # 外圈扩散：曼哈顿距离≤6 → 每轮+⌊异养-自养参数/3⌋
                energy_transfer = math.floor((neighbor.low_metabolism_high_metabolism - neighbor.autotrophic_heterotrophic) / 3)
            else:
                energy_transfer = 0
            
            if cell.energy > energy_transfer:
                cell.energy -= energy_transfer
                neighbor.energy += energy_transfer

    def _cell_movement(self, x, y):
        """
        细胞移动逻辑。
        :param x: 细胞的x坐标。
        :param y: 细胞的y坐标。
        """
        cell = self.grid[x][y]
        if not cell:
            return

        directions = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]
        random.shuffle(directions)
        for _ in range(2):  # 修改：增加细胞移动概率，尝试两次移动机会
            for dx, dy in directions:
                if cell.move(x, y, dx, dy):
                    break

    def get_average_energy(self):
        """
        计算网格中所有存活细胞的平均能量。
        :return: 平均能量值
        """
        total_energy = torch.sum(self.cell_data[:, :, CellDataIndex.ENERGY.value])
        alive_count = torch.sum(self.cell_data[:, :, CellDataIndex.ALIVE.value])
        return (total_energy / alive_count).item() if alive_count > 0 else 0

    def _energy_calculation(self):
        """
        向量化代谢和毒素计算。
        更新细胞能量和毒素状态。
        """
        traits_tensor = torch.stack([
            self.cell_data[:, :, CellDataIndex.LOW_METABOLISM_HIGH_METABOLISM.value],
            self.cell_data[:, :, CellDataIndex.MUTATION_RATE.value],
            self.cell_data[:, :, CellDataIndex.TOXIN_SECRETION.value],
            self.cell_data[:, :, CellDataIndex.TOXIN_RESISTANCE.value],
            self.cell_data[:, :, CellDataIndex.AUTOTROPHIC_HETEROTROPHIC.value],  # 新增：自养标志
            self.cell_data[:, :, CellDataIndex.DEFENSE_OFFENSE.value],  # 新增：防御/进攻参数
            self.cell_data[:, :, CellDataIndex.MOVE_ENERGY_COST.value],  # 新增：移动能量消耗
            self.cell_data[:, :, CellDataIndex.STATIC_MOTION.value]  # 新增：移动/静止
        ], dim=2)

        # 向量化代谢计算，考虑自养/异养和能量来源类型
        autotrophic_mask = (traits_tensor[:, :, 4] > 0) & self.cell_data[:, :, CellDataIndex.ALIVE.value]
        heterotrophic_mask = (traits_tensor[:, :, 4] <= 0) & self.cell_data[:, :, CellDataIndex.ALIVE.value]
        
        # 自养细胞能量增加
        energy_gain = traits_tensor[:, :, 4] * autotrophic_mask
        # 异养细胞能量消耗
        energy_delta = traits_tensor[:, :, 0] * (1 + traits_tensor[:, :, 1] * 10) * heterotrophic_mask
        
        self.cell_data[:, :, CellDataIndex.ENERGY.value] += energy_gain - energy_delta

        # 向量化毒素计算
        toxin_mask = (traits_tensor[:, :, 2] > 0) & self.cell_data[:, :, CellDataIndex.ALIVE.value]
        self.cell_data[:, :, CellDataIndex.TOXIN.value] += traits_tensor[:, :, 2] / 10 * toxin_mask
        self.cell_data[:, :, CellDataIndex.ENERGY.value] -= (1 + traits_tensor[:, :, 1] * 5) * toxin_mask

    def batch_update(self):
        """
        批量更新所有细胞的状态。
        包括年龄、存活时间、毒素扩散、能量计算和基因统计。
        """
        # 创建细胞状态掩码
        alive_mask = (self.cell_data[:, :, CellDataIndex.ALIVE.value] == 1)

        # 更新年龄和代数
        self.cell_data[:, :, CellDataIndex.AGE.value] += alive_mask
        # 更新存活时间
        self.cell_data[:, :, CellDataIndex.LIFETIME.value] += alive_mask
        # 同步更新Cell对象的lifetime
        for i in range(self.height):
            for j in range(self.width):
                if self.grid[i][j]:
                    self.grid[i][j].lifetime = self.cell_data[i, j, CellDataIndex.LIFETIME.value].item()
                    self.grid[i][j].age = self.cell_data[i, j, CellDataIndex.AGE.value].item()

        # 毒素扩散机制
        toxin_diffusion = self._toxin_diffusion()  # @call _toxin_diffusion
        # 应用扩散并考虑抗性
        self.cell_data[:, :, CellDataIndex.TOXIN.value] = torch.clamp(
            self.cell_data[:, :, CellDataIndex.TOXIN.value] * 0.7 + toxin_diffusion * self.cell_data[:, :, CellDataIndex.TOXIN_RESISTANCE.value],
            min=0,
            max=100
        )

        # 同步更新周围细胞的毒素值
        for i in range(self.height):
            for j in range(self.width):
                if self.grid[i][j]:
                    self.grid[i][j].toxin = self.cell_data[i, j, CellDataIndex.TOXIN.value].item()

        # 检查8细胞3代死亡规则
        kernel = torch.tensor([[1, 1, 1], [1, 0, 1], [1, 1, 1]], dtype=torch.float32, device=self.device)
        neighbor_counts = torch.conv2d(self.cell_data[:, :, CellDataIndex.ALIVE.value].unsqueeze(0).unsqueeze(0), kernel.unsqueeze(0).unsqueeze(0), padding=1).squeeze(0).squeeze(0)
        surrounded_mask = (neighbor_counts >= 8) & alive_mask

        # 能量计算
        self._energy_calculation()  # @call _energy_calculation

        # 同步更新Cell对象
        for i in range(self.height):
            for j in range(self.width):
                if self.grid[i][j]:
                    self.grid[i][j].energy = self.cell_data[i, j, CellDataIndex.ENERGY.value].item()
                    self.grid[i][j].age = self.cell_data[i, j, CellDataIndex.AGE.value].item()
                    self.grid[i][j].toxin = self.cell_data[i, j, CellDataIndex.TOXIN.value].item()
                    self.grid[i][j].lifetime = self.cell_data[i, j, CellDataIndex.LIFETIME.value].item()

        self._update_gene_stats()  # @call _update_gene_stats
        self.total_generations += 1